"""
LZW Compression Algorithm

Rubik3 Module
"""

class LZW:
    def encode(msg):
        """
        This function takes a string message, and returns a list of integers
        as encoded message.

        :param msg: The message to encode.
        :type msg: str
        :return: The encoded message.
        :rtype: list
        """

        # Initialize the dictionary with ASCII characters
        dic_size = 256 
        dic = {chr(i): i for i in range(dic_size)}

        P = ""
        code = []
        for C in msg:
            # Concatenate the current character with the previous string
            S = P + C
            # If the concatenated string is in the dictionary, update P to the new string
            if S in dic:
                P = S
            else:
                # If the string is not in the dictionary, append the code of P to the output
                # and add the new string to the dictionary
                code.append(dic[P])
                dic[S] = dic_size
                dic_size += 1
                P = C
        # If there's remaining string, append its code to the output
        if P:
            code.append(dic[P])
        return code

    def decode(code):
        """
        This function takes a list of integers as encoded message, and returns 
        the original string message.

        :param code: The encoded message.
        :type code: list
        :return: The original message.
        :rtype: str
        """

        # Initialize the dictionary with ASCII characters
        dic_size = 256
        dic = {i: chr(i) for i in range(dic_size)}

        # Initialize the output string with the first character of the code
        Output = c = chr(code.pop(0))

        for n in code:
            # If the code is in the dictionary, get the corresponding string
            if n in dic:
                courant = dic[n]
            # If the code is not in the dictionary, create a new string
            elif n == dic_size:
                courant = c + c[0]
            # Add the new string to the output
            Output += courant
            # Add the new string to the dictionary
            dic[dic_size] = c + courant[0]
            dic_size += 1
            c = courant
        return Output